import numpy as np 
import pandas as pd 
import scipy.sparse as sp

import torch.utils.data as data

import path_config


# 有些用户没有交互任何评论
def load_all():
	""" We load all the three file here to save time in each epoch. """
	train_data = pd.read_csv(path_config.train_path, sep=',', header=0, names=["author", "link_id", "lv0_id", "counts", "label"])
	test_data = pd.read_csv(path_config.test_path, sep=',', header=0, names=["author", "link_id", "lv0_id", "counts", "label"])

	user_num = train_data["author"].max() + 1
	comment_num = train_data["lv0_id"].max() + 1

	# Load interaction data as a dok matrix.
	train_mat = sp.dok_matrix((user_num, comment_num), dtype=np.int64)
	for i in range(user_num):
		user_id = train_data.iloc[i]["author"]
		comment_id = train_data.iloc[i]["lv0_id"]
		train_mat[user_id, comment_id] = 1

	test_mat = sp.dok_matrix((user_num, comment_num), dtype=np.int64)
	for i in range(user_num):
		user_id = test_data.iloc[i]["author"]
		comment_id = test_data.iloc[i]["lv0_id"]
		test_mat[user_id, comment_id] = 1

	return train_data, test_data, user_num, comment_num, train_mat


class BPRData(data.Dataset):
	def __init__(self, features,
				 item_num, train_mat=None, num_ng=0, is_training=None):
		super(BPRData, self).__init__()
		""" Note that the labels are only useful when training, we thus 
			add them in the ng_sample() function.
		"""
		self.features = features
		self.item_num = item_num
		self.train_mat = train_mat
		self.num_ng = num_ng
		self.is_training = is_training

	def ng_sample(self):
		assert self.is_training, 'no need to sampling when testing'

		self.features_fill = []
		for x in self.features:
			u, i = x[0], x[1]
			for t in range(self.num_ng):
				j = np.random.randint(self.item_num)
				while (u, j) in self.train_mat:
					j = np.random.randint(self.item_num)
				self.features_fill.append([u, i, j])

	def __len__(self):
		return self.num_ng * len(self.features) if self.is_training else len(self.features)

	def __getitem__(self, idx):
		features = self.features_fill if self.is_training else self.features

		user = features[idx][0]
		item_i = features[idx][1]
		item_j = features[idx][2] if self.is_training else features[idx][1]
		return user, item_i, item_j 


if __name__ == '__main__':
	load_all()